import os
from datetime import datetime
import dgl
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import torch

__all__ = [
    'Logger', 'print_system_info', 'visualize_graph', 'cosine_similarity', 'merge_nodes_v2'
]


class Logger:
    _instance = None

    def __new__(cls, name, level):
        if cls._instance is None:
            cls._instance = super(Logger, cls).__new__(cls)
            cls._instance.name = name
            cls._instance.level = level
        return cls._instance

    def info(self, message):
        return f"[INFO] {self.name}: {message}"

    def error(self, message):
        return f"[ERROR] {self.name}: {message}"


def print_system_info():
    print(f"PyTorch版本: {torch.__version__}")
    print(f"CUDA可用性: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"当前设备: {torch.cuda.get_device_name(0)}")
    print(f"DGL版本: {dgl.__version__}")
    print(np.__version__)
    print(f"Matplotlib版本: {matplotlib.__version__}")


def visualize_graph(g, feature_names=None):
    pos = g.ndata['pos'].detach().numpy()
    src, dst = g.edges()
    src = src.numpy()
    dst = dst.numpy()

    plt.figure(figsize=(8, 6))
    plt.scatter(pos[:, 0], pos[:, 1], c='blue')

    show_node_feat = []
    show_edge_feat = []
    if feature_names and isinstance(feature_names, (list, tuple)):
        for name in feature_names:
            if name in g.ndata:
                show_node_feat.append(name)
            if name in g.edata:
                show_edge_feat.append(name)

    for i, (x, y) in enumerate(pos):
        label = f'{i}'
        if show_node_feat:
            for feat in show_node_feat:
                value = g.ndata[feat][i].detach().cpu().numpy()
                label += f'\n{feat}:{value}'
        plt.text(x, y, label, fontsize=9, ha='right')

    for idx, (s, d) in enumerate(zip(src, dst)):
        x1, y1 = pos[s]
        x2, y2 = pos[d]
        plt.plot([x1, x2], [y1, y2], 'gray', alpha=0.5)
        if show_edge_feat:
            label = ''
            for feat in show_edge_feat:
                value = g.edata[feat][idx].detach().cpu().numpy()
                label += f'{feat}:{value} '
            if label:
                mx, my = (x1 + x2) / 2, (y1 + y2) / 2
                plt.text(mx, my, label.strip(), color='red', fontsize=8)

    plt.title('Graph Visualization')
    plt.xlabel('X')
    plt.ylabel('Y')

    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    save_folder = os.path.join(project_root, 'save')
    date_str = datetime.now().strftime("%Y-%m-%d")
    date_folder = os.path.join(save_folder, date_str)
    timestamp = datetime.now().strftime("%H%M%S_%f")
    filename = f"graph_{timestamp}.png"
    save_path = os.path.join(date_folder, filename)
    os.makedirs(date_folder, exist_ok=True)
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    print(f"图像已保存到: {save_path}")
    plt.show()


def cosine_similarity(a, b):
    a_norm = torch.norm(a, dim=1)
    b_norm = torch.norm(b, dim=1)
    dot_product = torch.sum(a * b, dim=1)
    sim = dot_product / (a_norm * b_norm)
    sim = torch.clamp(sim, min=0)
    return 0.1 * sim


# 新增: progress_idx / progress_total 用于外部传入当前合并序号与总次数；pbar 可传入一个 tqdm 实例
# 兼容：保持原有位置参数不变，新增参数在最后，旧代码无需修改即可不输出任何多余日志。
def merge_nodes_v2(g, node1_id, node2_id, edge_feat_key='dist', progress_idx=None, progress_total=None, pbar=None):
    # 获取两个节点的特征，取平均值作为node2的新特征
    node1_feat = g.ndata['feat'][node1_id]
    node2_feat = g.ndata['feat'][node2_id]
    avg_feat = (node1_feat + node2_feat) / 2.0
    g.ndata['feat'][node2_id] = avg_feat
    g.ndata['goods'][node2_id] += g.ndata['goods'][node1_id]
    # ==============aw


    # ==============aw
    id1 = g.ndata['id'][node1_id]
    id2 = g.ndata['id'][node2_id]
    fus_id = id1 | id2
    g.ndata['id'][node2_id] = fus_id
    src, dst = g.edges()
    src = src.tolist()
    dst = dst.tolist()
    num_nodes = g.num_nodes()
    edge_feats = {k: v.clone() for k, v in g.edata.items()}
    keep = []
    for i, (s, d) in enumerate(zip(src, dst)):
        if (s == node1_id and d == node2_id) or (s == node2_id and d == node1_id):
            continue
        keep.append(i)
    src = [src[i] for i in keep]
    dst = [dst[i] for i in keep]
    for k in edge_feats:
        edge_feats[k] = edge_feats[k][keep]
    for i in range(len(src)):
        if src[i] == node1_id:
            src[i] = node2_id
        if dst[i] == node1_id:
            dst[i] = node2_id
    keep_nodes = [i for i in range(num_nodes) if i != node1_id]
    node_map = {old: new for new, old in enumerate(keep_nodes)}
    new_node_feats = {k: v[keep_nodes] for k, v in g.ndata.items()}
    new_src = [node_map[s] for s in src]
    new_dst = [node_map[d] for d in dst]
    edge_dict = {}
    for i, (s, d) in enumerate(zip(new_src, new_dst)):
        key = (s, d)
        val = edge_feats[edge_feat_key][i].item()
        if key not in edge_dict or val < edge_dict[key][0]:
            edge_dict[key] = (val, i)
    keep_edge_idx = [idx for _, idx in edge_dict.values()]
    g_new = dgl.DGLGraph()
    g_new.add_nodes(len(keep_nodes))
    g_new.add_edges([new_src[i] for i in keep_edge_idx], [new_dst[i] for i in keep_edge_idx])
    for k in new_node_feats:
        g_new.ndata[k] = new_node_feats[k]
    for k in edge_feats:
        g_new.edata[k] = edge_feats[k][keep_edge_idx]

    # 进度展示：若传入 tqdm 实例则 update；否则如果给了 progress_idx/total 就用文本进度条。
    if pbar is not None:
        try:
            pbar.update(1)
        except Exception:
            pass
    elif progress_idx is not None and progress_total is not None and progress_total > 0:
        ratio = min(max(progress_idx / progress_total, 0), 1)
        bar_len = 30
        filled = int(bar_len * ratio)
        bar = '█' * filled + '-' * (bar_len - filled)
        end_char = '\n' if progress_idx >= progress_total else ''
        print(f"\r节点合并进度 |{bar}| {progress_idx}/{progress_total} ({ratio:6.2%})", end=end_char)

    return g_new
